package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
	"strings"

	mqtt "github.com/eclipse/paho.mqtt.golang"
	"github.com/gin-gonic/gin"
	"github.com/redis/go-redis/v9"
	"golang.org/x/net/context"
)

const (
    RedisHost          = "localhost"  
    RedisPort          = 6379         
    BrokerAddress      = "localhost"  
    BrokerPort         = 1883         
    MqttRequestTopic   = "request"
    MqttResponseTopic  = "response/#"
    ServerPort         = 55010
    RequestTimeout     = 120 * time.Second
    StaleCleanupPeriod = 60 * time.Second
    StaleTimeout       = 300 * time.Second
)

type CANMessage struct {
    CANID       uint32  `json:"can_id"`
    Timestamp   int64   `json:"timestamp"`
    Data        []interface{}  `json:"data"`
    VehicleID   string  `json:"vehicle_id"`
    IsAnomaly   bool    `json:"is_anomaly"`
    Confidence  float64 `json:"confidence,omitempty"`
}

type AnalysisResult struct {
    CANID       uint32            `json:"can_id"`
    VehicleID   string            `json:"vehicle_id"`
    AttackType  int               `json:"attack_type"`
    AttackName  string            `json:"attack_name"`
    Confidence  float64           `json:"confidence"`
    Model       string            `json:"model"`
    Timestamp   int64             `json:"timestamp"`
}

type ResponseLock struct {
	Event        chan struct{}
	ResponseData map[string]interface{}
	CreationTime time.Time
}

var (
	mqttClient      mqtt.Client
	redisClient     *redis.Client
	responseLocks   = make(map[string]*ResponseLock)
	locksMutex      sync.Mutex
	ctx             = context.Background()
	mqttInitialized = make(chan struct{})
)

func setupRedisClient() error {
	redisAddr := fmt.Sprintf("%s:%d", RedisHost, RedisPort)
	redisClient = redis.NewClient(&redis.Options{
		Addr:     redisAddr,
		Password: "", 
		DB:       0,  
	})

	_, err := redisClient.Ping(ctx).Result()
	if err != nil {
		return fmt.Errorf("Failed to connect to Redis: %v", err)
	}
	log.Println("Successfully connected to Redis")
	return nil
}

func setupMQTTClient() error {
	opts := mqtt.NewClientOptions()
	brokerURL := fmt.Sprintf("tcp://%s:%d", BrokerAddress, BrokerPort)
	opts.AddBroker(brokerURL)
	opts.SetClientID("GoBackend")
	opts.SetCleanSession(false)
	opts.SetKeepAlive(60 * time.Second)
	opts.SetAutoReconnect(true)
	opts.SetMaxReconnectInterval(60 * time.Second)

	opts.SetOnConnectHandler(func(client mqtt.Client) {
		log.Println("Connected to MQTT broker")
		if token := client.Subscribe(MqttResponseTopic, 0, onMessageReceived); token.Wait() && token.Error() != nil {
			log.Printf("Subscription failed: %v", token.Error())
		} else {
			log.Printf("Subscribed to %s", MqttResponseTopic)
		}
		close(mqttInitialized)
	})

	opts.SetConnectionLostHandler(func(client mqtt.Client, err error) {
		log.Printf("Connection lost: %v", err)
	})

	mqttClient = mqtt.NewClient(opts)
	if token := mqttClient.Connect(); token.Wait() && token.Error() != nil {
		return fmt.Errorf("Failed to connect to MQTT broker: %v", token.Error())
	}
	
	<-mqttInitialized
	
	return nil
}

func onMessageReceived(client mqtt.Client, msg mqtt.Message) {
	log.Printf("Received message on topic %s: %s", msg.Topic(), string(msg.Payload()))

	uuid := ""
	topicParts := strings.Split(msg.Topic(), "/")
	if len(topicParts) >= 2 {
		uuid = topicParts[1]
	} else {
		var payload map[string]interface{}
		if err := json.Unmarshal(msg.Payload(), &payload); err == nil {
			if userID, ok := payload["user_id"].(string); ok {
				uuid = userID
			}
		}
	}

	if uuid == "" {
		log.Println("Unable to extract UUID from message")
		return
	}

	var responseData map[string]interface{}
	if err := json.Unmarshal(msg.Payload(), &responseData); err != nil {
		log.Printf("Failed to parse response data: %v", err)
		return
	}

	locksMutex.Lock()
	lock, exists := responseLocks[uuid]
	locksMutex.Unlock()

	if exists {
		lock.ResponseData = responseData
		close(lock.Event)
	} else {
		jsonData, _ := json.Marshal(responseData)
		err := redisClient.Set(ctx, fmt.Sprintf("response:%s", uuid), string(jsonData), 5*time.Minute).Err()
		if err != nil {
			log.Printf("Failed to store response in Redis: %v", err)
		} else {
			log.Printf("%s has no waiting goroutine, stored in Redis", uuid)
		}
	}
}

func publishMessage(uuid string) error {
	if !mqttClient.IsConnected() {
		return fmt.Errorf("MQTT client is not connected")
	}

	token := mqttClient.Publish(MqttRequestTopic, 0, false, uuid)
	if token.Wait() && token.Error() != nil {
		return token.Error()
	}
	log.Printf("Published %s to %s", uuid, MqttRequestTopic)
	return nil
}

func saveToRedis(uuid string, data interface{}, maxRetries int) error {
	var err error
	for i := 0; i < maxRetries; i++ {
		jsonData, err := json.Marshal(data)
		if err != nil {
			return fmt.Errorf("JSON encoding failed: %v", err)
		}

		err = redisClient.Set(ctx, uuid, string(jsonData), 30*time.Minute).Err()
		if err == nil {
			log.Printf("Successfully saved %s to Redis", uuid)
			return nil
		}

		log.Printf("Redis error: %v. Retrying %d/%d...", err, i+1, maxRetries)
		time.Sleep(1 * time.Second)
	}
	return err
}

func cleanupStaleLocks() {
	for {
		time.Sleep(StaleCleanupPeriod)
		now := time.Now()
		
		locksMutex.Lock()
		var expiredUUIDs []string
		for uuid, lock := range responseLocks {
			if now.Sub(lock.CreationTime) > StaleTimeout {
				expiredUUIDs = append(expiredUUIDs, uuid)
			}
		}
		
		for _, uuid := range expiredUUIDs {
			delete(responseLocks, uuid)
		}
		locksMutex.Unlock()
		
		if len(expiredUUIDs) > 0 {
			log.Printf("Cleaned up %d expired locks", len(expiredUUIDs))
		}
	}
}

func handleCANAnomaly(c *gin.Context) {
	var canMessage CANMessage
	
	if err := c.ShouldBindJSON(&canMessage); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid CAN message format"})
		return
	}
	
	uuid := fmt.Sprintf("vehicle:%s:can:%d:%d", canMessage.VehicleID, canMessage.CANID, canMessage.Timestamp)
	log.Printf("Handling anomaly CAN message, UUID: %s", uuid)
	
	if err := saveToRedis(uuid, canMessage, 5); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to save CAN data to Redis"})
		return
	}

	lock := &ResponseLock{
		Event:        make(chan struct{}),
		ResponseData: nil,
		CreationTime: time.Now(),
	}
	locksMutex.Lock()
	responseLocks[uuid] = lock
	locksMutex.Unlock()

	if err := publishMessage(uuid); err != nil {
		locksMutex.Lock()
		delete(responseLocks, uuid)
		locksMutex.Unlock()
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to publish message to MQTT"})
		return
	}

	log.Printf("Waiting for response for %s...", uuid)

	select {
	case <-lock.Event:
		responseData := lock.ResponseData
		
		locksMutex.Lock()
		delete(responseLocks, uuid)
		locksMutex.Unlock()
		
		log.Printf("Received response for %s: %v", uuid, responseData)
		c.JSON(http.StatusOK, responseData)
		return
		
	case <-time.After(RequestTimeout):
		locksMutex.Lock()
		delete(responseLocks, uuid)
		locksMutex.Unlock()
		
		c.JSON(http.StatusGatewayTimeout, gin.H{"error": "Timed out waiting for worker node response"})
		return
	}
}

func handleRequest(c *gin.Context) {
	var data map[string]interface{}

	if err := c.BindJSON(&data); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid JSON"})
		return
	}
	uid, ok := data["user_id"].(string)
	if !ok {
		uid = "unknown"
	}
	uuid := fmt.Sprintf("user:%s", uid)
	log.Printf("Handling request, UUID: %s", uuid)

	if err := saveToRedis(uuid, data, 5); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to save data to Redis"})
		return
	}

	lock := &ResponseLock{
		Event:        make(chan struct{}),
		ResponseData: nil,
		CreationTime: time.Now(),
	}
	locksMutex.Lock()
	responseLocks[uuid] = lock
	locksMutex.Unlock()

	if err := publishMessage(uuid); err != nil {
		locksMutex.Lock()
		delete(responseLocks, uuid)
		locksMutex.Unlock()
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to publish message to MQTT"})
		return
	}

	log.Printf("Waiting for response for %s...", uuid)

	select {
	case <-lock.Event:
		responseData := lock.ResponseData
		
		locksMutex.Lock()
		delete(responseLocks, uuid)
		locksMutex.Unlock()
		
		log.Printf("Received response for %s: %v", uuid, responseData)
		c.JSON(http.StatusOK, responseData)
		return
		
	case <-time.After(RequestTimeout):
		locksMutex.Lock()
		delete(responseLocks, uuid)
		locksMutex.Unlock()
		
		c.JSON(http.StatusGatewayTimeout, gin.H{"error": "Timed out waiting for worker node response"})
		return
	}
}

func main() {
	if err := setupRedisClient(); err != nil {
		log.Fatalf("Failed to set up Redis: %v", err)
	}

	mqttInitialized = make(chan struct{})
	if err := setupMQTTClient(); err != nil {
		log.Fatalf("Failed to set up MQTT: %v", err)
	}

	go cleanupStaleLocks()

	router := gin.Default()
	router.POST("/test-api", handleRequest)         
	router.POST("/api/can/anomalies", handleCANAnomaly) 

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	
	go func() {
		<-quit
		log.Println("Shutting down server...")
		
		if mqttClient != nil && mqttClient.IsConnected() {
			mqttClient.Disconnect(250)
			log.Println("MQTT client disconnected")
		}
		
		if redisClient != nil {
			if err := redisClient.Close(); err != nil {
				log.Printf("Error closing Redis client: %v", err)
			} else {
				log.Println("Redis client closed")
			}
		}
		
		os.Exit(0)
	}()

	serverAddr := fmt.Sprintf("0.0.0.0:%d", ServerPort)
	log.Printf("Starting server on %s", serverAddr)
	if err := router.Run(serverAddr); err != nil {
		log.Fatalf("Failed to start server: %v", err)
	}
}
